import pandas as pd

def calculate_subsidies():

    student_info = pd.read_csv('data/2_student_info.csv')
    student_info = student_info[student_info['bf_zhusu'] == 1.0]

    consumption = pd.read_csv('data/7_consumption.csv')
    consumption['MonDeal'] = consumption['MonDeal'].abs()
    
    consumption_sum = consumption.groupby('bf_StudentID')['MonDeal'].sum().reset_index(name='total_consumption')
    
    merged_data = pd.merge(
        student_info,
        consumption_sum,
        on='bf_StudentID',
        how='left'
    ).fillna({'total_consumption': 0})

    threshold = merged_data['total_consumption'].quantile(0.2)
    poor_students = merged_data[merged_data['total_consumption'] <= threshold]

    median_consump = merged_data['total_consumption'].median()
    poor_students['津贴金额'] = (median_consump - poor_students['total_consumption']) * 0.6

    sorted_students = poor_students.sort_values('津贴金额', ascending=False)

    return [
        {
            "学号": row['bf_StudentID'],
            "姓名": row['bf_Name'],
            "性别": row['bf_sex'],
            "出生日期": str(int(row['bf_BornDate'])) if not pd.isna(row['bf_BornDate']) else '',
            "津贴金额": round(max(row['津贴金额'], 0), 2)
        }
        for _, row in sorted_students.iterrows()
    ]

def generate_subsidies():
    result = calculate_subsidies()
    df = pd.DataFrame(result)
    df.to_csv('data/subsidies.csv', index=False, encoding='utf-8-sig')

if __name__ == '__main__':
    generate_subsidies()
    # result = calculate_subsidy()
    # from rich import print
    # print(len(result))
    # print(result)
